/**
* This file is part of ORB-SLAM.
*
* Copyright (C) 2014 Raúl Mur-Artal <raulmur at unizar dot es> (University of Zaragoza)
* For more information see <http://webdiis.unizar.es/~raulmur/orbslam/>
*
* ORB-SLAM is free software: you can redistribute it and/or modify
* it under the terms of the GNU General Public License as published by
* the Free Software Foundation, either version 3 of the License, or
* (at your option) any later version.
*
* ORB-SLAM is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with ORB-SLAM. If not, see <http://www.gnu.org/licenses/>.
*/

#include "LoopClosing.hpp"
#include "Sim3Solver.hpp"
#include "Converter.hpp"
#include "Optimizer.hpp"
#include "ORBmatcher.hpp"
#include "Thirdparty/g2o/g2o/types/sim3/types_seven_dof_expmap.h"

#include <boost/date_time/posix_time/posix_time.hpp>
#include <boost/thread/thread.hpp>

namespace ORB_SLAM
{
    LoopClosing::LoopClosing() {}
    
    LoopClosing::LoopClosing(Map *pMap, KeyFrameDatabase *pDB, ORBVocabulary *pVoc):
    mbResetRequested(false), mpMap(pMap), mpKeyFrameDB(pDB), mpORBVocabulary(pVoc), mLastLoopKFid(0)
    {
        mnCovisibilityConsistencyTh = 3;
        mpMatchedKF = NULL;
    }
    
    void LoopClosing::SetTracker(Tracking *pTracker)
    {
        mpTracker=pTracker;
    }
    
    void LoopClosing::SetLocalMapper(LocalMapping *pLocalMapper)
    {
        mpLocalMapper=pLocalMapper;
    }
    
    
    void LoopClosing::Run()
    {
        
        //ros::Rate r(200);
        
        //while(ros::ok())
        while(true)
        {
            // Check if there are keyframes in the queue
            if(CheckNewKeyFrames())
            {
                // Detect loop candidates and check covisibility consistency
                if(DetectLoop())
                {
                    // Compute similarity transformation [sR|t]
                    if(ComputeSim3())
                    {
                        // Perform loop fusion and pose graph optimization
                        CorrectLoop();
                    }
                }
            }
            
            ResetIfRequested();
            //r.sleep();
            boost::this_thread::sleep(boost::posix_time::milliseconds(200));
        }
    }
    
    void LoopClosing::InsertKeyFrame(KeyFrame *pKF)
    {
        boost::mutex::scoped_lock lock(mMutexLoopQueue);
        if(pKF->mnId!=0)
            mlpLoopKeyFrameQueue.push_back(pKF);
    }
    
    bool LoopClosing::CheckNewKeyFrames()
    {
        boost::mutex::scoped_lock lock(mMutexLoopQueue);
        return(!mlpLoopKeyFrameQueue.empty());
    }
    
    bool LoopClosing::DetectLoop()
    {
        {
            boost::mutex::scoped_lock lock(mMutexLoopQueue);
            mpCurrentKF = mlpLoopKeyFrameQueue.front();
            mlpLoopKeyFrameQueue.pop_front();
            // Avoid that a keyframe can be erased while it is being process by this thread
            mpCurrentKF->SetNotErase();
        }
        
        //If the map contains less than 10 KF or less than 10KF have passed from last loop detection
        if(mpCurrentKF->mnId<mLastLoopKFid+10)
        {
            mpKeyFrameDB->add(mpCurrentKF);
            mpCurrentKF->SetErase();
            return false;
        }
        
        // Compute reference BoW similarity score
        // This is the lowest score to a connected keyframe in the covisibility graph
        // We will impose loop candidates to have a higher similarity than this
        vector<KeyFrame*> vpConnectedKeyFrames = mpCurrentKF->GetVectorCovisibleKeyFrames();
        DBoW2::BowVector CurrentBowVec = mpCurrentKF->GetBowVector();
        float minScore = 1;
        for(size_t i=0; i<vpConnectedKeyFrames.size(); i++)
        {
            KeyFrame* pKF = vpConnectedKeyFrames[i];
            if(pKF->isBad())
                continue;
            DBoW2::BowVector BowVec = pKF->GetBowVector();
            
            float score = mpORBVocabulary->score(CurrentBowVec, BowVec);
            
            if(score<minScore)
                minScore = score;
        }
        
        // Query the database imposing the minimum score
        vector<KeyFrame*> vpCandidateKFs = mpKeyFrameDB->DetectLoopCandidates(mpCurrentKF, minScore);
        
        
        // If there are no loop candidates, just add new keyframe and return false
        if(vpCandidateKFs.empty())
        {
            mpKeyFrameDB->add(mpCurrentKF);
            mvConsistentGroups.clear();
            mpCurrentKF->SetErase();
            return false;
        }
        
        // For each loop candidate check consistency with previous loop candidates
        // Each candidate expands a covisibility group (keyframes connected to the loop candidate in the covisibility graph)
        // A group is consistent with a previous group if they share at least a keyframe
        // We must detect a consistent loop in several consecutive keyframe to accept it
        mvpEnoughConsistentCandidates.clear();
        
        vector<ConsistentGroup> vCurrentConsistentGroups;
        vector<bool> vbConsistentGroup(mvConsistentGroups.size(),false);
        for(size_t i=0, iend=vpCandidateKFs.size(); i<iend; i++)
        {
            KeyFrame* pCandidateKF = vpCandidateKFs[i];
            
            set<KeyFrame*> spCandidateGroup = pCandidateKF->GetConnectedKeyFrames();
            spCandidateGroup.insert(pCandidateKF);
            
            bool bEnoughConsistent = false;
            bool bConsistentForSomeGroup = false;
            for(size_t iG=0, iendG=mvConsistentGroups.size(); iG<iendG; iG++)
            {
                set<KeyFrame*> sPreviousGroup = mvConsistentGroups[iG].first;
                
                bool bConsistent = false;
                for(set<KeyFrame*>::iterator sit=spCandidateGroup.begin(), send=spCandidateGroup.end(); sit!=send;sit++)
                {
                    if(sPreviousGroup.count(*sit))
                    {
                        bConsistent=true;
                        bConsistentForSomeGroup=true;
                        break;
                    }
                }
                
                if(bConsistent)
                {
                    int nPreviousConsistency = mvConsistentGroups[iG].second;
                    int nCurrentConsistency = nPreviousConsistency + 1;
                    if(!vbConsistentGroup[iG])
                    {
                        ConsistentGroup cg = make_pair(spCandidateGroup,nCurrentConsistency);
                        vCurrentConsistentGroups.push_back(cg);
                        vbConsistentGroup[iG]=true; //this avoid to include the same group more than once
                    }
                    if(nCurrentConsistency>=mnCovisibilityConsistencyTh && !bEnoughConsistent)
                    {
                        mvpEnoughConsistentCandidates.push_back(pCandidateKF);
                        bEnoughConsistent=true; //this avoid to insert the same candidate more than once
                    }
                }
            }
            
            // If the group is not consistent with any previous group insert with consistency counter set to zero
            if(!bConsistentForSomeGroup)
            {
                ConsistentGroup cg = make_pair(spCandidateGroup,0);
                vCurrentConsistentGroups.push_back(cg);
            }
        }
        
        // Update Covisibility Consistent Groups
        mvConsistentGroups = vCurrentConsistentGroups;
        
        
        // Add Current Keyframe to database
        mpKeyFrameDB->add(mpCurrentKF);
        
        if(mvpEnoughConsistentCandidates.empty())
        {
            mpCurrentKF->SetErase();
            return false;
        }
        else
        {
            return true;
        }
        
        mpCurrentKF->SetErase();
        return false;
    }
    
    bool LoopClosing::ComputeSim3()
    {
        // For each consistent loop candidate we try to compute a Sim3
        
        const int nInitialCandidates = mvpEnoughConsistentCandidates.size();
        
        // We compute first ORB matches for each candidate
        // If enough matches are found, we setup a Sim3Solver
        ORBmatcher matcher(0.75,true);
        
        vector<Sim3Solver*> vpSim3Solvers;
        vpSim3Solvers.resize(nInitialCandidates);
        
        vector<vector<MapPoint*> > vvpMapPointMatches;
        vvpMapPointMatches.resize(nInitialCandidates);
        
        vector<bool> vbDiscarded;
        vbDiscarded.resize(nInitialCandidates);
        
        int nCandidates=0; //candidates with enough matches
        
        for(int i=0; i<nInitialCandidates; i++)
        {
            KeyFrame* pKF = mvpEnoughConsistentCandidates[i];
            
            // avoid that local mapping erase it while it is being processed in this thread
            pKF->SetNotErase();
            
            if(pKF->isBad())
            {
                vbDiscarded[i] = true;
                continue;
            }
            
            int nmatches = matcher.SearchByBoW(mpCurrentKF,pKF,vvpMapPointMatches[i]);
            
            if(nmatches<20)
            {
                vbDiscarded[i] = true;
                continue;
            }
            else
            {
                Sim3Solver* pSolver = new Sim3Solver(mpCurrentKF,pKF,vvpMapPointMatches[i]);
                pSolver->SetRansacParameters(0.99,20,300);
                vpSim3Solvers[i] = pSolver;
            }
            
            nCandidates++;
        }
        
        bool bMatch = false;
        
        // Perform alternatively RANSAC iterations for each candidate
        // until one is succesful or all fail
        while(nCandidates>0 && !bMatch)
        {
            for(int i=0; i<nInitialCandidates; i++)
            {
                if(vbDiscarded[i])
                    continue;
                
                KeyFrame* pKF = mvpEnoughConsistentCandidates[i];
                
                // Perform 5 Ransac Iterations
                vector<bool> vbInliers;
                int nInliers;
                bool bNoMore;
                
                Sim3Solver* pSolver = vpSim3Solvers[i];
                cv::Mat Scm  = pSolver->iterate(5,bNoMore,vbInliers,nInliers);
                
                // If Ransac reachs max. iterations discard keyframe
                if(bNoMore)
                {
                    vbDiscarded[i]=true;
                    nCandidates--;
                }
                
                // If RANSAC returns a Sim3, perform a guided matching and optimize with all correspondences
                if(!Scm.empty())
                {
                    vector<MapPoint*> vpMapPointMatches(vvpMapPointMatches[i].size(), static_cast<MapPoint*>(NULL));
                    for(size_t j=0, jend=vbInliers.size(); j<jend; j++)
                    {
                        if(vbInliers[j])
                            vpMapPointMatches[j]=vvpMapPointMatches[i][j];
                    }
                    
                    cv::Mat R = pSolver->GetEstimatedRotation();
                    cv::Mat t = pSolver->GetEstimatedTranslation();
                    const float s = pSolver->GetEstimatedScale();
                    matcher.SearchBySim3(mpCurrentKF,pKF,vpMapPointMatches,s,R,t,7.5);
                    
                    
                    g2o::Sim3 gScm(Converter::toMatrix3d(R),Converter::toVector3d(t),s);
                    const int nInliers = Optimizer::OptimizeSim3(mpCurrentKF, pKF, vpMapPointMatches, gScm, 10);
                    
                    // If optimization is succesful stop ransacs and continue
                    if(nInliers>=20)
                    {
                        bMatch = true;
                        mpMatchedKF = pKF;
                        g2o::Sim3 gSmw(Converter::toMatrix3d(pKF->GetRotation()),Converter::toVector3d(pKF->GetTranslation()),1.0);
                        mg2oScw = gScm*gSmw;
                        mScw = Converter::toCvMat(mg2oScw);
                        
                        mvpCurrentMatchedPoints = vpMapPointMatches;
                        break;
                    }
                }
            }
        }
        
        if(!bMatch)
        {
            for(int i=0; i<nInitialCandidates; i++)
                mvpEnoughConsistentCandidates[i]->SetErase();
            mpCurrentKF->SetErase();
            return false;
        }
        
        // Retrieve MapPoints seen in Loop Keyframe and neighbors
        vector<KeyFrame*> vpLoopConnectedKFs = mpMatchedKF->GetVectorCovisibleKeyFrames();
        vpLoopConnectedKFs.push_back(mpMatchedKF);
        mvpLoopMapPoints.clear();
        for(vector<KeyFrame*>::iterator vit=vpLoopConnectedKFs.begin(); vit!=vpLoopConnectedKFs.end(); vit++)
        {
            KeyFrame* pKF = *vit;
            vector<MapPoint*> vpMapPoints = pKF->GetMapPointMatches();
            for(size_t i=0, iend=vpMapPoints.size(); i<iend; i++)
            {
                MapPoint* pMP = vpMapPoints[i];
                if(pMP)
                {
                    if(!pMP->isBad() && pMP->mnLoopPointForKF!=mpCurrentKF->mnId)
                    {
                        mvpLoopMapPoints.push_back(pMP);
                        pMP->mnLoopPointForKF=mpCurrentKF->mnId;
                    }
                }
            }
        }
        
        // Find more matches projecting with the computed Sim3
        matcher.SearchByProjection(mpCurrentKF, mScw, mvpLoopMapPoints, mvpCurrentMatchedPoints,10);
        
        // If enough matches accept Loop
        int nTotalMatches = 0;
        for(size_t i=0; i<mvpCurrentMatchedPoints.size(); i++)
        {
            if(mvpCurrentMatchedPoints[i])
                nTotalMatches++;
        }
        
        if(nTotalMatches>=40)
        {
            for(int i=0; i<nInitialCandidates; i++)
                if(mvpEnoughConsistentCandidates[i]!=mpMatchedKF)
                    mvpEnoughConsistentCandidates[i]->SetErase();
            return true;
        }
        else
        {
            for(int i=0; i<nInitialCandidates; i++)
                mvpEnoughConsistentCandidates[i]->SetErase();
            mpCurrentKF->SetErase();
            return false;
        }
        
    }
    
    void LoopClosing::CorrectLoop()
    {
        // Send a stop signal to Local Mapping
        // Avoid new keyframes are inserted while correcting the loop
        mpLocalMapper->RequestStop();
        
        // Wait until Local Mapping has effectively stopped
        //ros::Rate r(1e4);
        //while(ros::ok() && !mpLocalMapper->isStopped())
        while(!mpLocalMapper->isStopped())
        {
            //r.sleep();
            boost::this_thread::sleep(boost::posix_time::milliseconds(10000));
        }
        
        // Ensure current keyframe is updated
        mpCurrentKF->UpdateConnections();
        
        // Retrive keyframes connected to the current keyframe and compute corrected Sim3 pose by propagation
        mvpCurrentConnectedKFs = mpCurrentKF->GetVectorCovisibleKeyFrames();
        mvpCurrentConnectedKFs.push_back(mpCurrentKF);
        
        KeyFrameAndPose CorrectedSim3, NonCorrectedSim3;
        CorrectedSim3[mpCurrentKF]=mg2oScw;
        cv::Mat Twc = mpCurrentKF->GetPoseInverse();
        
        
        for(vector<KeyFrame*>::iterator vit=mvpCurrentConnectedKFs.begin(), vend=mvpCurrentConnectedKFs.end(); vit!=vend; vit++)
        {
            KeyFrame* pKFi = *vit;
            
            cv::Mat Tiw = pKFi->GetPose();
            
            if(pKFi!=mpCurrentKF)
            {
                cv::Mat Tic = Tiw*Twc;
                cv::Mat Ric = Tic.rowRange(0,3).colRange(0,3);
                cv::Mat tic = Tic.rowRange(0,3).col(3);
                g2o::Sim3 g2oSic(Converter::toMatrix3d(Ric),Converter::toVector3d(tic),1.0);
                g2o::Sim3 g2oCorrectedSiw = g2oSic*mg2oScw;
                //Pose corrected with the Sim3 of the loop closure
                CorrectedSim3[pKFi]=g2oCorrectedSiw;
            }
            
            cv::Mat Riw = Tiw.rowRange(0,3).colRange(0,3);
            cv::Mat tiw = Tiw.rowRange(0,3).col(3);
            g2o::Sim3 g2oSiw(Converter::toMatrix3d(Riw),Converter::toVector3d(tiw),1.0);
            //Pose without correction
            NonCorrectedSim3[pKFi]=g2oSiw;
        }
        
        // Correct all MapPoints obsrved by current keyframe and neighbors, so that they align with the other side of the loop
        for(KeyFrameAndPose::iterator mit=CorrectedSim3.begin(), mend=CorrectedSim3.end(); mit!=mend; mit++)
        {
            KeyFrame* pKFi = mit->first;
            g2o::Sim3 g2oCorrectedSiw = mit->second;
            g2o::Sim3 g2oCorrectedSwi = g2oCorrectedSiw.inverse();
            
            g2o::Sim3 g2oSiw =NonCorrectedSim3[pKFi];
            
            vector<MapPoint*> vpMPsi = pKFi->GetMapPointMatches();
            for(size_t iMP=0, endMPi = vpMPsi.size(); iMP<endMPi; iMP++)
            {
                MapPoint* pMPi = vpMPsi[iMP];
                if(!pMPi)
                    continue;
                if(pMPi->isBad())
                    continue;
                if(pMPi->mnCorrectedByKF==mpCurrentKF->mnId)
                    continue;
                
                // Project with non-corrected pose and project back with corrected pose
                cv::Mat P3Dw = pMPi->GetWorldPos();
                Eigen::Matrix<double,3,1> eigP3Dw = Converter::toVector3d(P3Dw);
                Eigen::Matrix<double,3,1> eigCorrectedP3Dw = g2oCorrectedSwi.map(g2oSiw.map(eigP3Dw));
                
                cv::Mat cvCorrectedP3Dw = Converter::toCvMat(eigCorrectedP3Dw);
                pMPi->SetWorldPos(cvCorrectedP3Dw);
                pMPi->mnCorrectedByKF = mpCurrentKF->mnId;
                pMPi->mnCorrectedReference = pKFi->mnId;
                pMPi->UpdateNormalAndDepth();
            }
            
            // Update keyframe pose with corrected Sim3. First transform Sim3 to SE3 (scale translation)
            Eigen::Matrix3d eigR = g2oCorrectedSiw.rotation().toRotationMatrix();
            Eigen::Vector3d eigt = g2oCorrectedSiw.translation();
            double s = g2oCorrectedSiw.scale();
            
            eigt *=(1./s); //[R t/s;0 1]
            
            cv::Mat correctedTiw = Converter::toCvSE3(eigR,eigt);
            
            pKFi->SetPose(correctedTiw);
            
            // Make sure connections are updated
            pKFi->UpdateConnections();
        }
        
        // Start Loop Fusion
        // Update matched map points and replace if duplicated
        for(size_t i=0; i<mvpCurrentMatchedPoints.size(); i++)
        {
            if(mvpCurrentMatchedPoints[i])
            {
                MapPoint* pLoopMP = mvpCurrentMatchedPoints[i];
                MapPoint* pCurMP = mpCurrentKF->GetMapPoint(i);
                if(pCurMP)
                    pCurMP->Replace(pLoopMP);
                else
                {
                    mpCurrentKF->AddMapPoint(pLoopMP,i);
                    pLoopMP->AddObservation(mpCurrentKF,i);
                    pLoopMP->ComputeDistinctiveDescriptors();
                }
            }
        }
        
        // Project MapPoints observed in the neighborhood of the loop keyframe
        // into the current keyframe and neighbors using corrected poses.
        // Fuse duplications.
        SearchAndFuse(CorrectedSim3);
        
        
        // After the MapPoint fusion, new links in the covisibility graph will appear attaching both sides of the loop
        map<KeyFrame*, set<KeyFrame*> > LoopConnections;
        
        for(vector<KeyFrame*>::iterator vit=mvpCurrentConnectedKFs.begin(), vend=mvpCurrentConnectedKFs.end(); vit!=vend; vit++)
        {
            KeyFrame* pKFi = *vit;
            vector<KeyFrame*> vpPreviousNeighbors = pKFi->GetVectorCovisibleKeyFrames();
            
            // Update connections. Detect new links.
            pKFi->UpdateConnections();
            LoopConnections[pKFi]=pKFi->GetConnectedKeyFrames();
            for(vector<KeyFrame*>::iterator vit_prev=vpPreviousNeighbors.begin(), vend_prev=vpPreviousNeighbors.end(); vit_prev!=vend_prev; vit_prev++)
            {
                LoopConnections[pKFi].erase(*vit_prev);
            }
            for(vector<KeyFrame*>::iterator vit2=mvpCurrentConnectedKFs.begin(), vend2=mvpCurrentConnectedKFs.end(); vit2!=vend2; vit2++)
            {
                LoopConnections[pKFi].erase(*vit2);
            }
        }
        
        mpTracker->ForceRelocalisation();
        
        Optimizer::OptimizeEssentialGraph(mpMap, mpMatchedKF, mpCurrentKF,  mg2oScw, NonCorrectedSim3, CorrectedSim3, LoopConnections);
        
        //Add edge
        mpMatchedKF->AddLoopEdge(mpCurrentKF);
        mpCurrentKF->AddLoopEdge(mpMatchedKF);
        
        std::cout << "Loop Closed!" << std::endl;
        
        // Loop closed. Release Local Mapping.
        mpLocalMapper->Release();
        
        mpMap->SetFlagAfterBA();
        
        mLastLoopKFid = mpCurrentKF->mnId;
    }
    
    void LoopClosing::SearchAndFuse(KeyFrameAndPose &CorrectedPosesMap)
    {
        ORBmatcher matcher(0.8);
        
        for(KeyFrameAndPose::iterator mit=CorrectedPosesMap.begin(), mend=CorrectedPosesMap.end(); mit!=mend;mit++)
        {
            KeyFrame* pKF = mit->first;
            
            g2o::Sim3 g2oScw = mit->second;
            cv::Mat cvScw = Converter::toCvMat(g2oScw);
            
            matcher.Fuse(pKF,cvScw,mvpLoopMapPoints,4);
        }
    }
    
    
    void LoopClosing::RequestReset()
    {
        {
            boost::mutex::scoped_lock lock(mMutexReset);
            mbResetRequested = true;
        }
        
        //ros::Rate r(500);
        //while(ros::ok())
        while(true)
        {
            {
                boost::mutex::scoped_lock lock2(mMutexReset);
                if(!mbResetRequested)
                    break;
            }
            //r.sleep();
            boost::this_thread::sleep(boost::posix_time::milliseconds(500));
        }
    }
    
    void LoopClosing::ResetIfRequested()
    {
        boost::mutex::scoped_lock lock(mMutexReset);
        if(mbResetRequested)
        {
            mlpLoopKeyFrameQueue.clear();
            mLastLoopKFid=0;
            mbResetRequested=false;
        }
    }
    
} //namespace ORB_SLAM
